import torch

class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv = torch.nn.Sequential(
            # 卷积层：输入通道1，输出通道32，卷积核5，步长2
            torch.nn.Conv2d(1, 32, kernel_size=5, stride=2),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)  # 池化层：窗口2x2
        )

        # 手动计算：输入 28x28 -> Conv+Pool 后是 6x6，通道数为32
        self.fc = torch.nn.Linear(32 * 6 * 6, 10)

    def forward(self, x):
        out = self.conv(x)
        out = out.view(out.size(0), -1)  # 展平
        out = self.fc(out)
        return out
